from flask import Blueprint, request, jsonify
from models.gpt2_model import get_model  # 导入模型
from product.production import Production  # 导入产品类

# 创建蓝图
api_bp = Blueprint('api', __name__)

# 获取 GPT2 模型和分词器
model, tokenizer = get_model()

# 路由 1：根路由，返回产品信息
@api_bp.route('/', methods=['GET'])
def hello_world():
    product = Production('Laptop', 1000, 3)
    return jsonify(product.to_dict())

# 路由 2：创建用户
@api_bp.route('/create_user', methods=['POST'])
def create_user():
    # 获取请求中的 JSON 数据
    data = request.get_json()

    if not data:
        return jsonify({"error": "Invalid JSON"}), 400  # 如果没有接收到有效的 JSON 数据，返回错误信息

    # 从 JSON 数据中提取字段
    user_id = data.get('user_id')
    username = data.get('username')
    email = data.get('email')

    # 返回 JSON 响应
    return jsonify({
        'message': 'User created successfully',
        'user': {
            'user_id': user_id,
            'username': username,
            'email': email
        }
    })

# 路由 3：生成文本
@api_bp.route('/generate', methods=['POST'])
def generate_text():
    # 从请求中获取输入文本
    data = request.get_json()
    prompt = data.get('prompt', '')  # 提示文本，默认空字符串

    if not prompt:
        return jsonify({"error": "No prompt provided"}), 400

    # 将输入文本转换为模型的输入格式
    inputs = tokenizer.encode(prompt, return_tensors="pt")

    # 使用模型生成文本
    output = model.generate(inputs, max_length=200, num_return_sequences=1, no_repeat_ngram_size=2)

    # 解码生成的文本
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

    # 返回生成的文本作为响应
    return jsonify({"generated_text": generated_text})
